#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
LoRA训练器常量定义
包含所有训练器相关的类型定义和常量
"""

# ============================================================================
# 训练器专用参数类型定义
# ============================================================================

# FLUX系列训练参数类型
FLUX_TRAINING_PARAMS = "FLUX_TRAINING_PARAMS"

# SD3系列训练参数类型
SD3_TRAINING_PARAMS = "SD3_TRAINING_PARAMS"

# SD3缓存参数类型
SD3_CACHE_PARAMS = "SD3_CACHE_PARAMS"

# SD3采样参数类型
SD3_SAMPLING_PARAMS = "SD3_SAMPLING_PARAMS"

# SD3分桶参数类型
SD3_BUCKET_PARAMS = "SD3_BUCKET_PARAMS"

# SDXL系列训练参数类型
SDXL_TRAINING_PARAMS = "SDXL_TRAINING_PARAMS"

# SD系列训练参数类型
SD_TRAINING_PARAMS = "TRAINING_PARAMS"

# 基础训练参数类型
BASIC_TRAINING_PARAMS = "BASIC_TRAINING_PARAMS"

# ============================================================================
# 内部连接类型定义（用于参数组合节点）
# ============================================================================

# FLUX内部连接类型
FLUX_DEPS_CHUNK = "FLUX_DEPS_CHUNK"
FLUX_TRAIN_CHUNK = "FLUX_TRAIN_CHUNK"
FLUX_PERF_CHUNK = "FLUX_PERF_CHUNK"

# ============================================================================
# 训练器状态常量
# ============================================================================

# 训练状态
TRAINING_STATUS_IDLE = "idle"
TRAINING_STATUS_RUNNING = "running"
TRAINING_STATUS_COMPLETED = "completed"
TRAINING_STATUS_FAILED = "failed"
TRAINING_STATUS_CANCELLED = "cancelled"

# ============================================================================
# 默认配置常量
# ============================================================================

# 默认输出目录
DEFAULT_OUTPUT_DIR = "./output"

# 默认模型保存格式
DEFAULT_SAVE_FORMAT = "safetensors"

# 默认随机种子
DEFAULT_SEED = 42

# ============================================================================
# 插件标识常量
# ============================================================================

# 统一的aux_id，符合GitHub格式
AUX_ID = "fengya/comfyui_lora_train"

# 插件信息
PLUGIN_NAME = "ComfyUI LoRA训练插件"
PLUGIN_VERSION = "1.0.0"
PLUGIN_AUTHOR = "comfyui_lora_train"

# ============================================================================
# 导出所有常量
# ============================================================================

__all__ = [
    # 训练参数类型
    "FLUX_TRAINING_PARAMS",
    "SD3_TRAINING_PARAMS", 
    "SD3_CACHE_PARAMS",
    "SD3_SAMPLING_PARAMS",
    "SD3_BUCKET_PARAMS",
    "SDXL_TRAINING_PARAMS",
    "SD_TRAINING_PARAMS",
    "BASIC_TRAINING_PARAMS",
    
    # 内部连接类型
    "FLUX_DEPS_CHUNK",
    "FLUX_TRAIN_CHUNK", 
    "FLUX_PERF_CHUNK",
    
    # 训练状态
    "TRAINING_STATUS_IDLE",
    "TRAINING_STATUS_RUNNING",
    "TRAINING_STATUS_COMPLETED",
    "TRAINING_STATUS_FAILED",
    "TRAINING_STATUS_CANCELLED",
    
    # 默认配置
    "DEFAULT_OUTPUT_DIR",
    "DEFAULT_SAVE_FORMAT",
    "DEFAULT_SEED",
    "DEFAULT_VAE_NAME",
    
    # 节点分类常量
    "SDXL_TRAINING_PARAMS_CATEGORY",
    "SDXL_LORA_TRAINER_CATEGORY",
    "SD3_TRAINING_PARAMS_CATEGORY",
    "SD3_LORA_TRAINER_CATEGORY",
    "FLUX_NEW_PARAMS_CATEGORY",
    "FLUX_TRAINING_CATEGORY",
    "FLUX_DEBUG_CATEGORY",
    "BASIC_LORA_TRAINING_CATEGORY",
    "ADVANCED_LORA_TRAINER_CATEGORY",
    "CHARACTER_LORA_PARAMS_CATEGORY",
    "LORA_PARAM_CONFIG_CATEGORY",
    "LORA_PERFORMANCE_CONFIG_CATEGORY",
    "LORA_CONFIG_MERGER_CATEGORY",
    "LORA_PARAMS_ORGANIZER_CATEGORY",
    "PERFORMANCE_CONFIG_CATEGORY",
]

# ============================================================================
# 节点分类常量
# ============================================================================

# SDXL相关节点分类
SDXL_TRAINING_PARAMS_CATEGORY = "AI训练/SDXL LoRA训练/参数"
SDXL_LORA_TRAINER_CATEGORY = "AI训练/SDXL LoRA训练"

# SD3相关节点分类
SD3_TRAINING_PARAMS_CATEGORY = "AI训练/SD3 LoRA训练/参数"
SD3_LORA_TRAINER_CATEGORY = "AI训练/SD3 LoRA训练"

# FLUX相关节点分类
FLUX_NEW_PARAMS_CATEGORY = "AI训练/FLUX训练/新参数"
FLUX_TRAINING_CATEGORY = "AI训练/FLUX训练"
FLUX_DEBUG_CATEGORY = "AI训练/FLUX调试"

# 通用LoRA训练节点分类
BASIC_LORA_TRAINING_CATEGORY = "AI训练/LoRA训练"
ADVANCED_LORA_TRAINER_CATEGORY = "AI训练/高级LoRA训练"
CHARACTER_LORA_PARAMS_CATEGORY = "AI训练/高级LoRA训练/参数"

# LoRA配置相关节点分类
LORA_PARAM_CONFIG_CATEGORY = "AI训练/LoRA配置"
LORA_PERFORMANCE_CONFIG_CATEGORY = "AI训练/测试/LoRA配置"
LORA_CONFIG_MERGER_CATEGORY = "AI训练/测试/LoRA配置"

# 测试相关节点分类
LORA_PARAMS_ORGANIZER_CATEGORY = "AI训练/测试/LoRA参数"
PERFORMANCE_CONFIG_CATEGORY = "AI训练/LoRA训练"

# VAE相关常量
DEFAULT_VAE_NAME = "无（使用模型自带VAE）"
